# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

import os
import pandas
import numpy
from absl import app, flags
from io import StringIO

from coin.strategy.mm.tool.feed_stat_util import enumerate_leading_digit
import coin.strategy.mm.tool.archive_base as abase

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages


kBenford = '''col	value
trade_leadingdigit_1_w	0.3010
trade_leadingdigit_2_w	0.1760
trade_leadingdigit_3_w	0.1250
trade_leadingdigit_4_w	0.0970
trade_leadingdigit_5_w	0.0790
trade_leadingdigit_6_w	0.0670
trade_leadingdigit_7_w	0.0580
trade_leadingdigit_8_w	0.0510
trade_leadingdigit_9_w	0.0450'''


def main(_):
  tdates = abase.get_trading_dates(flags.FLAGS.trading_date)
  os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
  benforddf = pandas.read_csv(StringIO(kBenford), sep='\t')
  benfordseri = benforddf['value']
  benfordseri.index = benforddf['col']

  for tdate in tdates:
    plt.rcParams['lines.linewidth'] = 0.5
    plt.rcParams['figure.figsize'] = 12, 8
    plt.rcParams['font.size'] = 10
    plt.rcParams['legend.fontsize'] = 10
    plt.rcParams['xtick.labelsize'] = 10
    plt.rcParams['ytick.labelsize'] = 10
    pandas.set_option("display.precision", 3)

    tdstr = tdate.strftime("%Y%m%d")
    legs = []
    plt.plot(benfordseri.iloc[1:], lw=1, markersize=2)
    legs.append('benford')
    tuples = []
    for nametuple, df in enumerate_leading_digit(tdstr, flags.FLAGS.pattern):
      tdigcols = [
          col for col in df.columns
          if col.find("trade_leadingdigit") >= 0
          and col.find("digit_0") < 0
          and col.find("digit_1") < 0]
      for _, tdgseri in df.iterrows():
        diff = (
            tdgseri[tdigcols] / tdgseri[tdigcols].sum()
            - benfordseri.iloc[1:] / benfordseri[1:].sum())
        rmse = "%.6f" % numpy.sqrt((diff * diff).mean())
        tuples.append((
            rmse,
            tdgseri[tdigcols],
            f"rmse_{rmse}_{tdgseri['exchange']}_{tdgseri['symbol']}"))
    for rmse, seri, legend in sorted(tuples):
      plt.plot(seri, lw=0.5, markersize=1)
      legs.append(legend)
    plt.title(f"{tdstr}_{flags.FLAGS.pattern}")
    plt.legend(legs, loc='upper right')
    plt.xticks(rotation=20)
    plt.savefig(f'{flags.FLAGS.out_dir}/{tdstr}_{flags.FLAGS.pattern}.png')
    plt.close()


if __name__ == '__main__':
  flags.DEFINE_string('pattern', '%', '')
  flags.DEFINE_string('out_dir', 'fake_volume', '')
  flags.DEFINE_string('trading_date', None, '')
  app.run(main)
